{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "qv_BkFZKteD1"
   },
   "source": [
    "# NER Powered Semantic Search"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "lpFX45v8tfsD"
   },
   "source": [
    "This notebook shows how to use Named Entity Recognition (NER) for vector search with LanceDB. We will:\n",
    "\n",
    "1. Extract named entities from text.\n",
    "2. Store them in a LanceDB as metadata (alongside respective text vectors).\n",
    "3. We extract named entities from incoming queries and use them to filter and search only through records containing these named entities.\n",
    "\n",
    "This is particularly helpful if you want to restrict the search score to records that contain information about the named entities that are also found within the query.\n",
    "\n",
    "Let's get started."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "wytJNKiIt7_w"
   },
   "source": [
    "# Installing Dependencies"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "zwhWLUUitU18",
    "outputId": "e3d0238f-643e-49d2-a73a-33ea059484f0"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m86.0/86.0 kB\u001b[0m \u001b[31m2.4 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
      "\u001b[?25h  Preparing metadata (setup.py) ... \u001b[?25l\u001b[?25hdone\n",
      "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m493.7/493.7 kB\u001b[0m \u001b[31m16.2 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
      "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m72.0/72.0 kB\u001b[0m \u001b[31m9.2 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
      "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m7.9/7.9 MB\u001b[0m \u001b[31m96.2 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
      "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m1.3/1.3 MB\u001b[0m \u001b[31m86.0 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
      "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m311.2/311.2 kB\u001b[0m \u001b[31m36.9 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
      "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m115.3/115.3 kB\u001b[0m \u001b[31m14.5 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
      "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m134.8/134.8 kB\u001b[0m \u001b[31m17.2 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
      "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m18.3/18.3 MB\u001b[0m \u001b[31m65.5 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
      "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m38.0/38.0 MB\u001b[0m \u001b[31m16.6 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
      "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m98.7/98.7 kB\u001b[0m \u001b[31m11.8 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
      "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m3.8/3.8 MB\u001b[0m \u001b[31m105.2 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
      "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m1.3/1.3 MB\u001b[0m \u001b[31m63.9 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
      "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m295.0/295.0 kB\u001b[0m \u001b[31m34.9 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
      "\u001b[?25h  Building wheel for sentence_transformers (setup.py) ... \u001b[?25l\u001b[?25hdone\n",
      "\u001b[31mERROR: pip's dependency resolver does not currently take into account all the packages that are installed. This behaviour is the source of the following dependency conflicts.\n",
      "ibis-framework 6.2.0 requires pyarrow<13,>=2, but you have pyarrow 14.0.0 which is incompatible.\n",
      "pandas-gbq 0.17.9 requires pyarrow<10.0dev,>=3.0.0, but you have pyarrow 14.0.0 which is incompatible.\u001b[0m\u001b[31m\n",
      "\u001b[0m"
     ]
    }
   ],
   "source": [
    "!pip install sentence_transformers datasets lancedb -qU"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "o674zy69uGiz"
   },
   "source": [
    "# Load and Prepare Datasets\n",
    "\n",
    "We use a dataset containing ~190K articles scraped from Medium. We select 50K articles from the dataset as indexing all the articles may take some time. This dataset can be loaded from the HuggingFace dataset hub as follows:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 212,
     "referenced_widgets": [
      "585a644ea5a24471a9d1fbee41a983f9",
      "b23dd6630d4843d39039ed884dca1001",
      "0acb4fb4b40f4381b625f6a080c6b864",
      "03aef87debd2499f8ad7d1f2e6419552",
      "9bfa0e084f134634abec0b8f8cf5e5c7",
      "264802d54f8049189d205103f6e24c14",
      "b65b30eac4294bd5b28641f3d8449fe6",
      "8b3eba38821a4a6e9d676c807b941376",
      "8fc5eb076e4e446e85ac946bc7fd89eb",
      "db5e14ee1ba44deda1c488240dfd843b",
      "720ac11a94c34e7faf977acf218a03f8",
      "d087948b6c594c1b853cc651978c263b",
      "47cd8a03f7004992b0642c7f14eb5aec",
      "d19548c9373f490d867d8d6eab13b53b",
      "e1cd874df82740a0b4ea11354f011cff",
      "b71c62ea02784d61aee39c1711f7a22f",
      "0634b14776054ef7836192afadd4942a",
      "fe674568b1d649d0a3e8b34f267e668b",
      "8b2e8c7450604d058487040e55cda058",
      "080c101dc9164df1a917d911491d3a67",
      "e195980f5b854d8db09975ba58920d68",
      "2f2f7e70e33e4e6092b1a45cbe383b29",
      "78fbab968c5642ad975937afd5597542",
      "c4b94c8ef29e44b9aa09f7c3d7ac8963",
      "b39f012194014a2fa2e5cdc09fa56169",
      "8a5caf0e4bdf474a8802c5c1408cd0a5",
      "08f350dfd08d4a18bd8fc17742f8be5e",
      "69928b9e3659493faa4b03a4bc061209",
      "93889f479b254847abcd43e31bbeb825",
      "fc048c1979ed4acda52cc766928d510d",
      "29e5ffd2a6f1433b8e02efb89b4cb721",
      "dab283e57d29454c9921acfbb98f4929",
      "d889fcb80d00490d95d579938dfef206",
      "a50b0d3ea66a42dba0b86999bbb1bf78",
      "e55657d612da484dbfbe2617f576baf5",
      "e2b5e82a41414e548033b4ae4831b5a5",
      "d3fcb48013b84d4eb876def088d94b36",
      "2b0ba7f41922431894764654c6e9da39",
      "5fdaedb7a1d4467487c1bd7f602dc666",
      "b19dd3edf08e4f8cb44d57817af781a1",
      "0859e8db5ef14ed5814c7e165c2085c4",
      "3d6cd06243fe4870a1d7573bd5e4ab8a",
      "e047a131330e4ee990700050084a02e7",
      "08f7661d4a0e4de2831cde83efafd083",
      "efb414dad5e34f7bbc29569f9f73efc9",
      "a8837ba89300422fae4dca6c4472d2cb",
      "1cfaa06af3394f2d87f52133d494b7b1",
      "14f2073bf2324c649d2910caab3a46c3",
      "541be2496a9b4bad85770dbb4362ffca",
      "71d25cd77c504150bdda3f04add78978",
      "a93a2319c5e14747a43d42f0896cffb0",
      "f5d9ae440d4c40e39999daf51bb46d02",
      "1e60faa61ea245689055cc843dd45387",
      "d3664ba62fe74bf58509a58376dc87ca",
      "8fef047db7ce4d798a7e3313a8c35383"
     ]
    },
    "id": "OU4BaxRFuJb7",
    "outputId": "200a6054-2c9a-4bb3-ce02-93bfff92bcb1"
   },
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "585a644ea5a24471a9d1fbee41a983f9",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Downloading readme:   0%|          | 0.00/2.26k [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "d087948b6c594c1b853cc651978c263b",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Downloading data files:   0%|          | 0/1 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "78fbab968c5642ad975937afd5597542",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Downloading data:   0%|          | 0.00/1.04G [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "a50b0d3ea66a42dba0b86999bbb1bf78",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Extracting data files:   0%|          | 0/1 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "efb414dad5e34f7bbc29569f9f73efc9",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Generating train split: 0 examples [00:00, ? examples/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/usr/local/lib/python3.10/dist-packages/datasets/table.py:1421: FutureWarning: promote has been superseded by mode='default'.\n",
      "  table = cls._concat_blocks(blocks, axis=0)\n"
     ]
    }
   ],
   "source": [
    "from datasets import load_dataset\n",
    "\n",
    "# load the dataset and convert to pandas dataframe\n",
    "df = load_dataset(\n",
    "    \"fabiochiu/medium-articles\", data_files=\"medium_articles.csv\", split=\"train\"\n",
    ").to_pandas()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "8BkI1uEVuX9I"
   },
   "source": [
    "## Preprocessing on dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 293
    },
    "id": "lLl9qJeJuRbI",
    "outputId": "a557eff3-e455-4c56-832d-67a60505d6b6"
   },
   "outputs": [
    {
     "data": {
      "text/html": [
       "\n",
       "  <div id=\"df-ac876683-aaf5-4647-b4bf-beb34ff99005\" class=\"colab-df-container\">\n",
       "    <div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>title</th>\n",
       "      <th>text</th>\n",
       "      <th>url</th>\n",
       "      <th>authors</th>\n",
       "      <th>timestamp</th>\n",
       "      <th>tags</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>4172</th>\n",
       "      <td>How the Data Stole Christmas</td>\n",
       "      <td>by Anonymous\\n\\nThe door sprung open and our t...</td>\n",
       "      <td>https://medium.com/data-ops/how-the-data-stole...</td>\n",
       "      <td>[]</td>\n",
       "      <td>2019-12-24 13:22:33.143000+00:00</td>\n",
       "      <td>['Data Science', 'Big Data', 'Dataops', 'Analy...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>174868</th>\n",
       "      <td>Automating Light Switch using the ESP32 Board ...</td>\n",
       "      <td>A story about how I escaped the boring task th...</td>\n",
       "      <td>https://python.plainenglish.io/automating-ligh...</td>\n",
       "      <td>['Tomas Rasymas']</td>\n",
       "      <td>2021-09-14 07:20:52.342000+00:00</td>\n",
       "      <td>['Programming', 'Python', 'Software Developmen...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>100171</th>\n",
       "      <td>Keep Going Quotes Sayings for When Hope is Lost</td>\n",
       "      <td>It’s a very thrilling thing to achieve a goal....</td>\n",
       "      <td>https://medium.com/@yourselfquotes/keep-going-...</td>\n",
       "      <td>['Yourself Quotes']</td>\n",
       "      <td>2021-01-05 12:13:04.018000+00:00</td>\n",
       "      <td>['Quotes']</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>141757</th>\n",
       "      <td>When Will the Smoke Clear From Bay Area Skies?</td>\n",
       "      <td>Bay Area cities are contending with some of th...</td>\n",
       "      <td>https://thebolditalic.com/when-will-the-smoke-...</td>\n",
       "      <td>['Matt Charnock']</td>\n",
       "      <td>2020-09-15 22:38:33.924000+00:00</td>\n",
       "      <td>['Bay Area', 'San Francisco', 'California', 'W...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>183489</th>\n",
       "      <td>The ABC’s of Sustainability… easy as 1, 2, 3</td>\n",
       "      <td>By Julia DiPrete\\n\\n(according to the Jackson ...</td>\n",
       "      <td>https://medium.com/sipwines/the-abcs-of-sustai...</td>\n",
       "      <td>['Sip Wines']</td>\n",
       "      <td>2021-03-02 23:39:49.948000+00:00</td>\n",
       "      <td>['Wine Tasting', 'Sustainability', 'Wine']</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>\n",
       "    <div class=\"colab-df-buttons\">\n",
       "\n",
       "  <div class=\"colab-df-container\">\n",
       "    <button class=\"colab-df-convert\" onclick=\"convertToInteractive('df-ac876683-aaf5-4647-b4bf-beb34ff99005')\"\n",
       "            title=\"Convert this dataframe to an interactive table.\"\n",
       "            style=\"display:none;\">\n",
       "\n",
       "  <svg xmlns=\"http://www.w3.org/2000/svg\" height=\"24px\" viewBox=\"0 -960 960 960\">\n",
       "    <path d=\"M120-120v-720h720v720H120Zm60-500h600v-160H180v160Zm220 220h160v-160H400v160Zm0 220h160v-160H400v160ZM180-400h160v-160H180v160Zm440 0h160v-160H620v160ZM180-180h160v-160H180v160Zm440 0h160v-160H620v160Z\"/>\n",
       "  </svg>\n",
       "    </button>\n",
       "\n",
       "  <style>\n",
       "    .colab-df-container {\n",
       "      display:flex;\n",
       "      gap: 12px;\n",
       "    }\n",
       "\n",
       "    .colab-df-convert {\n",
       "      background-color: #E8F0FE;\n",
       "      border: none;\n",
       "      border-radius: 50%;\n",
       "      cursor: pointer;\n",
       "      display: none;\n",
       "      fill: #1967D2;\n",
       "      height: 32px;\n",
       "      padding: 0 0 0 0;\n",
       "      width: 32px;\n",
       "    }\n",
       "\n",
       "    .colab-df-convert:hover {\n",
       "      background-color: #E2EBFA;\n",
       "      box-shadow: 0px 1px 2px rgba(60, 64, 67, 0.3), 0px 1px 3px 1px rgba(60, 64, 67, 0.15);\n",
       "      fill: #174EA6;\n",
       "    }\n",
       "\n",
       "    .colab-df-buttons div {\n",
       "      margin-bottom: 4px;\n",
       "    }\n",
       "\n",
       "    [theme=dark] .colab-df-convert {\n",
       "      background-color: #3B4455;\n",
       "      fill: #D2E3FC;\n",
       "    }\n",
       "\n",
       "    [theme=dark] .colab-df-convert:hover {\n",
       "      background-color: #434B5C;\n",
       "      box-shadow: 0px 1px 3px 1px rgba(0, 0, 0, 0.15);\n",
       "      filter: drop-shadow(0px 1px 2px rgba(0, 0, 0, 0.3));\n",
       "      fill: #FFFFFF;\n",
       "    }\n",
       "  </style>\n",
       "\n",
       "    <script>\n",
       "      const buttonEl =\n",
       "        document.querySelector('#df-ac876683-aaf5-4647-b4bf-beb34ff99005 button.colab-df-convert');\n",
       "      buttonEl.style.display =\n",
       "        google.colab.kernel.accessAllowed ? 'block' : 'none';\n",
       "\n",
       "      async function convertToInteractive(key) {\n",
       "        const element = document.querySelector('#df-ac876683-aaf5-4647-b4bf-beb34ff99005');\n",
       "        const dataTable =\n",
       "          await google.colab.kernel.invokeFunction('convertToInteractive',\n",
       "                                                    [key], {});\n",
       "        if (!dataTable) return;\n",
       "\n",
       "        const docLinkHtml = 'Like what you see? Visit the ' +\n",
       "          '<a target=\"_blank\" href=https://colab.research.google.com/notebooks/data_table.ipynb>data table notebook</a>'\n",
       "          + ' to learn more about interactive tables.';\n",
       "        element.innerHTML = '';\n",
       "        dataTable['output_type'] = 'display_data';\n",
       "        await google.colab.output.renderOutput(dataTable, element);\n",
       "        const docLink = document.createElement('div');\n",
       "        docLink.innerHTML = docLinkHtml;\n",
       "        element.appendChild(docLink);\n",
       "      }\n",
       "    </script>\n",
       "  </div>\n",
       "\n",
       "\n",
       "<div id=\"df-9904742c-a2eb-41fa-8754-14aff2c1fc56\">\n",
       "  <button class=\"colab-df-quickchart\" onclick=\"quickchart('df-9904742c-a2eb-41fa-8754-14aff2c1fc56')\"\n",
       "            title=\"Suggest charts\"\n",
       "            style=\"display:none;\">\n",
       "\n",
       "<svg xmlns=\"http://www.w3.org/2000/svg\" height=\"24px\"viewBox=\"0 0 24 24\"\n",
       "     width=\"24px\">\n",
       "    <g>\n",
       "        <path d=\"M19 3H5c-1.1 0-2 .9-2 2v14c0 1.1.9 2 2 2h14c1.1 0 2-.9 2-2V5c0-1.1-.9-2-2-2zM9 17H7v-7h2v7zm4 0h-2V7h2v10zm4 0h-2v-4h2v4z\"/>\n",
       "    </g>\n",
       "</svg>\n",
       "  </button>\n",
       "\n",
       "<style>\n",
       "  .colab-df-quickchart {\n",
       "      --bg-color: #E8F0FE;\n",
       "      --fill-color: #1967D2;\n",
       "      --hover-bg-color: #E2EBFA;\n",
       "      --hover-fill-color: #174EA6;\n",
       "      --disabled-fill-color: #AAA;\n",
       "      --disabled-bg-color: #DDD;\n",
       "  }\n",
       "\n",
       "  [theme=dark] .colab-df-quickchart {\n",
       "      --bg-color: #3B4455;\n",
       "      --fill-color: #D2E3FC;\n",
       "      --hover-bg-color: #434B5C;\n",
       "      --hover-fill-color: #FFFFFF;\n",
       "      --disabled-bg-color: #3B4455;\n",
       "      --disabled-fill-color: #666;\n",
       "  }\n",
       "\n",
       "  .colab-df-quickchart {\n",
       "    background-color: var(--bg-color);\n",
       "    border: none;\n",
       "    border-radius: 50%;\n",
       "    cursor: pointer;\n",
       "    display: none;\n",
       "    fill: var(--fill-color);\n",
       "    height: 32px;\n",
       "    padding: 0;\n",
       "    width: 32px;\n",
       "  }\n",
       "\n",
       "  .colab-df-quickchart:hover {\n",
       "    background-color: var(--hover-bg-color);\n",
       "    box-shadow: 0 1px 2px rgba(60, 64, 67, 0.3), 0 1px 3px 1px rgba(60, 64, 67, 0.15);\n",
       "    fill: var(--button-hover-fill-color);\n",
       "  }\n",
       "\n",
       "  .colab-df-quickchart-complete:disabled,\n",
       "  .colab-df-quickchart-complete:disabled:hover {\n",
       "    background-color: var(--disabled-bg-color);\n",
       "    fill: var(--disabled-fill-color);\n",
       "    box-shadow: none;\n",
       "  }\n",
       "\n",
       "  .colab-df-spinner {\n",
       "    border: 2px solid var(--fill-color);\n",
       "    border-color: transparent;\n",
       "    border-bottom-color: var(--fill-color);\n",
       "    animation:\n",
       "      spin 1s steps(1) infinite;\n",
       "  }\n",
       "\n",
       "  @keyframes spin {\n",
       "    0% {\n",
       "      border-color: transparent;\n",
       "      border-bottom-color: var(--fill-color);\n",
       "      border-left-color: var(--fill-color);\n",
       "    }\n",
       "    20% {\n",
       "      border-color: transparent;\n",
       "      border-left-color: var(--fill-color);\n",
       "      border-top-color: var(--fill-color);\n",
       "    }\n",
       "    30% {\n",
       "      border-color: transparent;\n",
       "      border-left-color: var(--fill-color);\n",
       "      border-top-color: var(--fill-color);\n",
       "      border-right-color: var(--fill-color);\n",
       "    }\n",
       "    40% {\n",
       "      border-color: transparent;\n",
       "      border-right-color: var(--fill-color);\n",
       "      border-top-color: var(--fill-color);\n",
       "    }\n",
       "    60% {\n",
       "      border-color: transparent;\n",
       "      border-right-color: var(--fill-color);\n",
       "    }\n",
       "    80% {\n",
       "      border-color: transparent;\n",
       "      border-right-color: var(--fill-color);\n",
       "      border-bottom-color: var(--fill-color);\n",
       "    }\n",
       "    90% {\n",
       "      border-color: transparent;\n",
       "      border-bottom-color: var(--fill-color);\n",
       "    }\n",
       "  }\n",
       "</style>\n",
       "\n",
       "  <script>\n",
       "    async function quickchart(key) {\n",
       "      const quickchartButtonEl =\n",
       "        document.querySelector('#' + key + ' button');\n",
       "      quickchartButtonEl.disabled = true;  // To prevent multiple clicks.\n",
       "      quickchartButtonEl.classList.add('colab-df-spinner');\n",
       "      try {\n",
       "        const charts = await google.colab.kernel.invokeFunction(\n",
       "            'suggestCharts', [key], {});\n",
       "      } catch (error) {\n",
       "        console.error('Error during call to suggestCharts:', error);\n",
       "      }\n",
       "      quickchartButtonEl.classList.remove('colab-df-spinner');\n",
       "      quickchartButtonEl.classList.add('colab-df-quickchart-complete');\n",
       "    }\n",
       "    (() => {\n",
       "      let quickchartButtonEl =\n",
       "        document.querySelector('#df-9904742c-a2eb-41fa-8754-14aff2c1fc56 button');\n",
       "      quickchartButtonEl.style.display =\n",
       "        google.colab.kernel.accessAllowed ? 'block' : 'none';\n",
       "    })();\n",
       "  </script>\n",
       "</div>\n",
       "    </div>\n",
       "  </div>\n"
      ],
      "text/plain": [
       "                                                    title  \\\n",
       "4172                         How the Data Stole Christmas   \n",
       "174868  Automating Light Switch using the ESP32 Board ...   \n",
       "100171    Keep Going Quotes Sayings for When Hope is Lost   \n",
       "141757     When Will the Smoke Clear From Bay Area Skies?   \n",
       "183489       The ABC’s of Sustainability… easy as 1, 2, 3   \n",
       "\n",
       "                                                     text  \\\n",
       "4172    by Anonymous\\n\\nThe door sprung open and our t...   \n",
       "174868  A story about how I escaped the boring task th...   \n",
       "100171  It’s a very thrilling thing to achieve a goal....   \n",
       "141757  Bay Area cities are contending with some of th...   \n",
       "183489  By Julia DiPrete\\n\\n(according to the Jackson ...   \n",
       "\n",
       "                                                      url  \\\n",
       "4172    https://medium.com/data-ops/how-the-data-stole...   \n",
       "174868  https://python.plainenglish.io/automating-ligh...   \n",
       "100171  https://medium.com/@yourselfquotes/keep-going-...   \n",
       "141757  https://thebolditalic.com/when-will-the-smoke-...   \n",
       "183489  https://medium.com/sipwines/the-abcs-of-sustai...   \n",
       "\n",
       "                    authors                         timestamp  \\\n",
       "4172                     []  2019-12-24 13:22:33.143000+00:00   \n",
       "174868    ['Tomas Rasymas']  2021-09-14 07:20:52.342000+00:00   \n",
       "100171  ['Yourself Quotes']  2021-01-05 12:13:04.018000+00:00   \n",
       "141757    ['Matt Charnock']  2020-09-15 22:38:33.924000+00:00   \n",
       "183489        ['Sip Wines']  2021-03-02 23:39:49.948000+00:00   \n",
       "\n",
       "                                                     tags  \n",
       "4172    ['Data Science', 'Big Data', 'Dataops', 'Analy...  \n",
       "174868  ['Programming', 'Python', 'Software Developmen...  \n",
       "100171                                         ['Quotes']  \n",
       "141757  ['Bay Area', 'San Francisco', 'California', 'W...  \n",
       "183489         ['Wine Tasting', 'Sustainability', 'Wine']  "
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# drop empty rows and select 20k articles\n",
    "df = df.dropna().sample(20000, random_state=32)\n",
    "df.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "id": "tWzqtfPUuWUv"
   },
   "outputs": [],
   "source": [
    "# select first 1000 characters\n",
    "df[\"text\"] = df[\"text\"].str[:1000]\n",
    "# join article title and the text\n",
    "df[\"title_text\"] = df[\"title\"] + \". \" + df[\"text\"]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "Hbj6S26euhTO"
   },
   "source": [
    "## Initialize NER model\n",
    "\n",
    "To extract named entities, we will use a NER model finetuned on a BERT-base model. The model can be loaded from the HuggingFace model hub as follows:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "id": "gVy7_MWzuk_M"
   },
   "outputs": [],
   "source": [
    "import torch\n",
    "\n",
    "# set device to GPU if available\n",
    "device = torch.cuda.current_device() if torch.cuda.is_available() else None"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 281,
     "referenced_widgets": [
      "d9424b33ffc547d7b16b96e753d697b0",
      "8dd1938e92e743778db2b8db53afb3cc",
      "b1b8210e70ad4f4fad07e2e3afc14c8c",
      "852b3c3ac82d401e9e6b826889e92409",
      "f27c8638a40744409422bdbe3c3b99c9",
      "8a9a8bffa3fc4f35923721c8e3dcacdf",
      "bc8d5a7306cf492eb195c4e97145f732",
      "e287196168424deba2f02b78e7ee3937",
      "0833cb0a6c174dfa982ea9ba1e36ddc5",
      "7ed8a5c5650d4b7a8448cd51741d593f",
      "d615e739da0c44d7acce7c60560e4a84",
      "163a866ed11748cfb6050139952df1ef",
      "515848a31c6b45cbac1655e25aa01784",
      "30a742b47c464140831b5a94cd8d23f6",
      "561715cdf7ef4e159bcae9236f9f540f",
      "6e969e0edfb4446aa5fe90ee387723e7",
      "5e5239727f034d678cfdefbf2ab75d2e",
      "fcd70f15b02f43509097c30cf4ba9153",
      "d68ba2a80e8342d1abf51f8ee9b8350c",
      "60f92482c9764672be103b1b17d74347",
      "4a40d95d764e490f95342cc8399e3fb0",
      "488e859cb46c42fe9a7491e9f1ac1b71",
      "085363a78410479aac5dcdcea40a295e",
      "6d2cbb2a5da24c8ba5790baf2d22be4c",
      "115e5a4af7d34c9a9b6d2ee3e0044428",
      "eef1cff9d1b74f60a41eb606294b8314",
      "219a5071dc1f4fa0a939f596efe7d4c7",
      "10ce238f585f4122b3dd95ffeec56936",
      "701571462a3f46de97082fef78ec5631",
      "87438c5455c942a6a44d315d7ff5023b",
      "925b7243d889407cb91895ffe5785a8f",
      "6d5114bdf0d3442f8b63359f022235b2",
      "ca3a3ebb28824683aac84117deef7ac9",
      "2cbd32b10d3e43e1b5a36f484e805084",
      "3952afec7c784db687e1a771e031646c",
      "66408705b7e848368e248485b10e8e95",
      "0c164cabaedb4b07a498bc6c810daba9",
      "566d435d1eab4e9f86db2b167dba0b1e",
      "62cf76f786404a919d2deecb3c92a3ef",
      "0a305d65dc784f13888f8f4c90509699",
      "331f1223e36b4e24bcdb6e2f3a2a2681",
      "5d0a63af0c0b483883d3c8d99360e23d",
      "63403ae2d3284df2b4d72221da186a19",
      "17b110a7670e40a6b9c9485e8d235481",
      "2d93b0d890144adb85a3760e46c6211f",
      "ea52a7e4a6e644a1a71925f777f8c5fa",
      "26a295f18931424bac6c68446288d7a7",
      "8239f19e7dc944c297c4598142923558",
      "0613b8c72d0d4cf8bc68cd79cdc5ee74",
      "ef35563a5b4544b5a1286d3da5c67bf2",
      "47e3d7e89ace47fe9311b49241640459",
      "55d4510d22e045248bd6cc172035bce8",
      "2ef96d13edab4f12b3d1072a1c32c1f7",
      "98e313f21305483d8773c8dc53e04ec0",
      "f2e60cec34f144f29e81de973470dada",
      "aa9e1d72063e495c85c423d70f376e46",
      "1750cd652f044733a743d8e73c0f2b78",
      "78476dfaac18485abd59d41650741625",
      "b2512ee942ea4b6dac29aa7576067d45",
      "9afc9172389a47a6aa95f55c0802827e",
      "b9b3b8032b864c47a8facc9f8dda5f1c",
      "fda98e4ff02c417e8eba9d83e623f4fd",
      "f771675ff3e840fa8e9535cbe9b4ca7b",
      "923c87913f9346dc8d920e4a666fecc5",
      "e2ab9b4bdaa14c4e85d49c0c8d782604",
      "d3d136771b6944a4a13540a69dd9e0a6"
     ]
    },
    "id": "_fv5oCUXuokl",
    "outputId": "ff2e5eb6-6216-4d26-ab6d-b13988da37cb"
   },
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "d9424b33ffc547d7b16b96e753d697b0",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Downloading (…)okenizer_config.json:   0%|          | 0.00/59.0 [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "163a866ed11748cfb6050139952df1ef",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Downloading (…)lve/main/config.json:   0%|          | 0.00/829 [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "085363a78410479aac5dcdcea40a295e",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Downloading (…)solve/main/vocab.txt:   0%|          | 0.00/213k [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "2cbd32b10d3e43e1b5a36f484e805084",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Downloading (…)in/added_tokens.json:   0%|          | 0.00/2.00 [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "2d93b0d890144adb85a3760e46c6211f",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Downloading (…)cial_tokens_map.json:   0%|          | 0.00/112 [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "aa9e1d72063e495c85c423d70f376e46",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Downloading model.safetensors:   0%|          | 0.00/433M [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Some weights of the model checkpoint at dslim/bert-base-NER were not used when initializing BertForTokenClassification: ['bert.pooler.dense.bias', 'bert.pooler.dense.weight']\n",
      "- This IS expected if you are initializing BertForTokenClassification from the checkpoint of a model trained on another task or with another architecture (e.g. initializing a BertForSequenceClassification model from a BertForPreTraining model).\n",
      "- This IS NOT expected if you are initializing BertForTokenClassification from the checkpoint of a model that you expect to be exactly identical (initializing a BertForSequenceClassification model from a BertForSequenceClassification model).\n"
     ]
    }
   ],
   "source": [
    "from transformers import AutoTokenizer, AutoModelForTokenClassification\n",
    "from transformers import pipeline\n",
    "\n",
    "model_id = \"dslim/bert-base-NER\"\n",
    "\n",
    "# load the tokenizer from huggingface\n",
    "tokenizer = AutoTokenizer.from_pretrained(model_id)\n",
    "# load the NER model from huggingface\n",
    "model = AutoModelForTokenClassification.from_pretrained(model_id)\n",
    "# load the tokenizer and model into a NER pipeline\n",
    "nlp = pipeline(\n",
    "    \"ner\", model=model, tokenizer=tokenizer, aggregation_strategy=\"max\", device=device\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "gbjv53IVuw8a",
    "outputId": "178325d5-da30-4d6c-e521-98197b0bc189"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[{'entity_group': 'LOC',\n",
       "  'score': 0.99969244,\n",
       "  'word': 'London',\n",
       "  'start': 37,\n",
       "  'end': 43}]"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "text = \"What are the best Places to visit in London\"\n",
    "# use the NER pipeline to extract named entities from the text\n",
    "nlp(text)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "hkqWpnSevBi0"
   },
   "source": [
    "Our NER pipeline is working as expected and accurately extracting entities from the text."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "YeeFIg-YvFpH"
   },
   "source": [
    "## Initialize Retreiver\n",
    "\n",
    "A retriever model is used to embed passages (article title + first 1000 characters) and queries. It creates embeddings such that queries and passages with similar meanings are close in the vector space. We will use a sentence-transformer model as our retriever. The model can be loaded as follows:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 552,
     "referenced_widgets": [
      "0ef84d3ae1004cdfb30ab691d2e4644d",
      "b3ab1a1b0f674eeda1432dbc8735898d",
      "e6588677d8f244f3971345bad3b445a2",
      "4fe50a002efe4cea837dbea55495606e",
      "4c0c70fbcc6648fc84e23a1afeeab6dc",
      "9ac6f19c42b2487389f599a0433d8200",
      "b8fb5157f2f34a08ac3755264ef92f88",
      "f3377417a7424b4e832840a9829cfc40",
      "f200bbd1e4404129834991b6d771e0a4",
      "ced342b85d394d6382aa330eb973ee61",
      "b95bb134564442e9b9cb22cfbfbc23a2",
      "9dffce3aa0b74469be61d11276352ece",
      "87ce082e714f47939426f4d59e8900ac",
      "9b1f0acae3f542b598ee592154121cbe",
      "d68baed34e6f4301b4ca86a2b274b603",
      "98f368e28ed24f1db5da2faf70ebef42",
      "c36fa0823faf42e9917c0826aac33d30",
      "87116fea3fea4e26bed4c11427778774",
      "eeba63fabeea4c12b20c8db68d959510",
      "fa5dcc3cad0c4533af1bc0c4b270d9de",
      "43f0de7066254eac9d5f9b2bc337e195",
      "ba9bd10251794cfd83688ec94bb12e3d",
      "5a85862df6e3465fbd631ea59fab267f",
      "a8dbe7dfd3384670b1ef8e8943fb2708",
      "0452481ac0cd41769a5437faff2cc451",
      "6446e38305524805a4625b8c360abde9",
      "14206793020c4d24b2c6554ebccb26b9",
      "43912e14f3c343da9f437a527b1b0b0c",
      "be6e3c066f594fcf99f97a0063a68367",
      "00a2343e2b5444f6b9a07a2fe74402c4",
      "d8fd0f53e3db4078bb29a5f29ade73ea",
      "01e0b64b23444ed391fb90f2183002f7",
      "c2dfcc65dcd249859c75fa60168a367a",
      "5e37fbc06f7a4c2abf0f58c2f91c3e5e",
      "54a47ff157c7415290722a0b8a5b291e",
      "e4b48bd72ab944aea13d94119d1d069a",
      "bdf756de1b714c2bbc334a73e1715501",
      "89355cfa688445539abfc539bab5899b",
      "59ae5700345a4ef981977a4165912fc8",
      "869927a602334d2a95465c69b2670f21",
      "cc319ad428174051bafa641fda722aab",
      "52fb01e3983a4dcdaa0a6c73f8b5cd23",
      "99e4b64f9fbf4694be8cec266267b399",
      "983ed125a337428fa8f3e1df4d92f022",
      "17716766563c4675bf61ac598d497e41",
      "6513335828d34195923b9e5daca59ca1",
      "d70f64d4565f4f1483a5d26b3556b8a7",
      "4e67b376461b4197933e0db95d927619",
      "04f4ceaf17244996890336cd33b55416",
      "0bdf96d521584b42b2d6cb5713eaace6",
      "54d39161c66d4fb4843b7d3a0a38115d",
      "72640f60c88e41dab5c942c74fecd881",
      "63f95a3f91c3463d98e9150f63a73b00",
      "36cb31db01b4448083814e535a65ea36",
      "86d7d92cda9449a38b9eace8713a15bb",
      "8a2dd43c2dd34c109dad128dc2b4b227",
      "9fb98cd620b040c3b0d6e7c007cb36cd",
      "2cb937659db14651b5794ff6d0e050ad",
      "733b6609614e4410ab61bb6d639df616",
      "cdcadcbded2146adb30312d4cb94328d",
      "4d1dee2846f34244b34d1272678d7d8c",
      "cb1b86ecdbfb410db8ce69afb3e1cb8f",
      "d05763758588477d9bb6599ce934d047",
      "e91495b9c61c4003bfc64dea86e1c06b",
      "87fb764430de4ef6aadca4c787f797de",
      "74a84adaea55458f97b519dbb4f58386",
      "b7f567522385408698b6c88755ed395f",
      "e34266bd9d434a2584eba0a96212d0ff",
      "134ea9bfd3b54712bfa5cb54fc045862",
      "8f4aa502a7ab41e7a6c667a2ebc1e01d",
      "2c7361b6d38c4ee48d0f02b67bf16642",
      "4cfa374d058e4978ba58dbba97029dd7",
      "5a4e9109511b4d45a30e44821eef439e",
      "cb77588c1af4431cbd3805300742595d",
      "e6faf300d4db4875b9591df9c5841ed0",
      "22f52d4a114548c9884c81f62e7d2982",
      "938c5c07d88c438e85ae68b92b66d034",
      "fd4e1e6b0191425fb23dda02b2126ec9",
      "8b57c82704a84951b92983b94cb3f3b4",
      "54ac9e213d554d58bc66a74a118e1030",
      "c6e94779b46f42d68a6b522a3a13d2b0",
      "f71945f5a8cd42c2841bfa2e19d7e247",
      "959fd365d37c4fe1a778e2aa0fd55c0f",
      "da2fb1a89deb459580bd5a02feba21bb",
      "0d3223538e3842fc8d7d1d8ab75a5938",
      "ae6e95ad731a4e99a1d10c8364409abf",
      "461703a84ac441b9a87ee36a4abaf1af",
      "508f9ca668d74c63886fd159defaa009",
      "4fd0a9acb66c4e979d8a943e5fb65bb7",
      "69f5662393544580b99ed721b635826c",
      "5d450805c12f48c28f2fe04bf9aa9293",
      "048484846b8346bc9b8f81ec01dfaaf4",
      "3fbd1f074ed642868fa89fc0def4a7aa",
      "42057ca3393f4899a4ab083c219aff6e",
      "6aef92d3adca4326a34389b6c6b6216d",
      "7405adf526ef4742a89e206758c210bd",
      "246ac973a67d460d9f81e346f25adbca",
      "968705197b7e44318f8c706f5b51245b",
      "ad0d6a41ea664209b1bd109d57c2b322",
      "79231ec70284441aaa47b2acec2ad698",
      "7c78c59da7b0403db56a5e4d6d4266d2",
      "6df83ae6232a4361ba4eda92c6d1e346",
      "91adfc5fd4844693bda4dd4c4b741961",
      "bce02837c8f14e62bc00b564ffb49e80",
      "c2d714d83e7b4245b201b10b72542fd8",
      "653fbf259e8f446e925a399304979cbd",
      "5863902daf63474f9255ec6b3bdf26b7",
      "df35cb02f21243bb81cc326872d8b8ad",
      "5d6a5c721dbd42a78ccf5c45a63b50be",
      "e1cb2070ceda4f8a89b719ba81ec51b0",
      "02d72daba5b744fa9dfbc03381a76d61",
      "968d29b7bb964e7a87d42e5d92514d66",
      "c6223b68c4414ab9b1494208560df422",
      "9621e8cb5bbd4d7580fe9f56fa7a982b",
      "93d88c377f5a4991b5054e8b90e82f82",
      "a9386915f91444ba89238c2de6111e91",
      "a19cb92b27a94a70842f7530796fafa0",
      "14da948ba91249b79ad6e10597b0f762",
      "cec4d718acce43a8bb48cda3e47a3375",
      "f0298f5e2a2d4f938d259ea91349f33d",
      "365f93a962d34aae86f37a8dbbdd209f",
      "227ef762045b44c2b799f334ade1e5f8",
      "ce870cfe0b444c148b84c464a52c463c",
      "eef6b3b840d94f3ca4fc646ff20a8a39",
      "760b97b874794e31a079d3618469e3c6",
      "61d492cd1fe84cf8b695ce4cca184e90",
      "6155725be53848e8821131baf932f57f",
      "306318c6f770469fbb04f406dd9763c7",
      "510382c14cda4a74b698b7dd8839c036",
      "22212043de5748cd8f45f316eeba9a21",
      "9c1d8e1187eb4f72b8b0adbd4f32c885",
      "8bfe1afcfe224fe3b049e63c0b99ca79",
      "3cd1b8ca9781438fa4846303a99a8d77",
      "2d4d64e5e3d344de9dda557797aaf738",
      "d91e3c3dec6c4924b8a8575a9ef522b8",
      "73301d9e2cde4888ace95cadb1c25267",
      "36c4a5bbc5e146d29b7c74c982eafe00",
      "bc66b3a2ff59471b9e5974b8a5e3edc1",
      "bfe35b13112041b5b8bf4d9d6fae739d",
      "5dfc922ade414f75bb0679778e4faa9c",
      "2f83ff9e128848f8808cb47b4c3b3701",
      "68eac5caf30848ee99a9fcc5e5299a55",
      "ee25d22a9af24e5a980916da0d6f4413",
      "8fc8d664606d4191896b9bfd5b35b4a2",
      "a8b7b6c007af4801bb895c1abe071727",
      "54ef076d1fe0473787efdec1ee082cfb",
      "b35352946acd48fe960a5535b3820b8c",
      "fdc3b678825a4c3db17fb7e7d663bd3a",
      "e4850a426cfb4e6d8972e9a6ce69f39a",
      "8d35ed741fa743779af6993c8f9c3f50",
      "01d23f17984d4265acb6290b046bc5c1",
      "b0cd330896634a2ba314249f6157b302",
      "01888c5ad1b8487099c6b16c6dad2a7d",
      "26601177c0ea4d7a8858d49072755322"
     ]
    },
    "id": "49K7UX9_vCDR",
    "outputId": "8becb3f2-d0aa-4a7a-db66-85755a0d1b83"
   },
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "0ef84d3ae1004cdfb30ab691d2e4644d",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Downloading (…)e933c/.gitattributes:   0%|          | 0.00/737 [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "9dffce3aa0b74469be61d11276352ece",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Downloading (…)_Pooling/config.json:   0%|          | 0.00/190 [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "5a85862df6e3465fbd631ea59fab267f",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Downloading (…)cbe6ee933c/README.md:   0%|          | 0.00/9.85k [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "5e37fbc06f7a4c2abf0f58c2f91c3e5e",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Downloading (…)e6ee933c/config.json:   0%|          | 0.00/591 [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "17716766563c4675bf61ac598d497e41",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Downloading (…)ce_transformers.json:   0%|          | 0.00/116 [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "8a2dd43c2dd34c109dad128dc2b4b227",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Downloading (…)33c/data_config.json:   0%|          | 0.00/15.7k [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "b7f567522385408698b6c88755ed395f",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Downloading pytorch_model.bin:   0%|          | 0.00/438M [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "fd4e1e6b0191425fb23dda02b2126ec9",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Downloading (…)nce_bert_config.json:   0%|          | 0.00/53.0 [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "4fd0a9acb66c4e979d8a943e5fb65bb7",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Downloading (…)cial_tokens_map.json:   0%|          | 0.00/239 [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "79231ec70284441aaa47b2acec2ad698",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Downloading (…)e933c/tokenizer.json:   0%|          | 0.00/466k [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "02d72daba5b744fa9dfbc03381a76d61",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Downloading (…)okenizer_config.json:   0%|          | 0.00/383 [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "227ef762045b44c2b799f334ade1e5f8",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Downloading (…)933c/train_script.py:   0%|          | 0.00/13.2k [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "3cd1b8ca9781438fa4846303a99a8d77",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Downloading (…)cbe6ee933c/vocab.txt:   0%|          | 0.00/232k [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "8fc8d664606d4191896b9bfd5b35b4a2",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Downloading (…)6ee933c/modules.json:   0%|          | 0.00/349 [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "SentenceTransformer(\n",
       "  (0): Transformer({'max_seq_length': 128, 'do_lower_case': False}) with Transformer model: MPNetModel \n",
       "  (1): Pooling({'word_embedding_dimension': 768, 'pooling_mode_cls_token': False, 'pooling_mode_mean_tokens': True, 'pooling_mode_max_tokens': False, 'pooling_mode_mean_sqrt_len_tokens': False})\n",
       "  (2): Normalize()\n",
       ")"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from sentence_transformers import SentenceTransformer\n",
    "\n",
    "# load the model from huggingface\n",
    "retriever = SentenceTransformer(\n",
    "    \"flax-sentence-embeddings/all_datasets_v3_mpnet-base\", device=device\n",
    ")\n",
    "retriever"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "gITGXAu9vVMn"
   },
   "source": [
    "## Initialize LanceDB"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {
    "id": "ADk9TBhuvYzT"
   },
   "outputs": [],
   "source": [
    "import lancedb\n",
    "\n",
    "db = lancedb.connect(\"./.lancedb\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "r9USO-wyvjJ-"
   },
   "source": [
    "## Generate Embeddings and Insert\n",
    "\n",
    "We generate embeddings for the title_text column we created earlier. Alongside the embeddings, we also include the named entities in the index as metadata. Later we will apply a filter based on these named entities when executing queries.\n",
    "\n",
    "Let's first write a helper function to extract named entities from a batch of text."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {
    "id": "EsSGeJrCvo17"
   },
   "outputs": [],
   "source": [
    "def extract_named_entities(text_batch):\n",
    "    # extract named entities using the NER pipeline\n",
    "    extracted_batch = nlp(text_batch)\n",
    "    entities = []\n",
    "    # loop through the results and only select the entity names\n",
    "    for text in extracted_batch:\n",
    "        ne = [entity[\"word\"] for entity in text]\n",
    "        entities.append(ne)\n",
    "    return entities"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 49,
     "referenced_widgets": [
      "2660a58840304f55b81b61a775c0a5d4",
      "a5845cdb905d41fd8cab4f1545174182",
      "ae49f70e0d774e888ccf94367e379ef5",
      "b5ddef19b49f44788ae2dee56fe66873",
      "f312bbf7fb464c02bcf879cd4a3a3c04",
      "afd1d42394354cbe93372500b435deae",
      "243d811f59274eba94db0fde4a9d40fb",
      "f909372292e84f9badaf8730cc3d3229",
      "ca7339e987f84d359c29de1f6b093e8e",
      "b5579a51c4af46af9deb6fff571514a6",
      "fa11337b1ae54fd5a708eaacc9d73bfd"
     ]
    },
    "id": "1caMEcVFwOeO",
    "outputId": "8c68d2cf-9223-4336-8a56-069a4d96661a"
   },
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "2660a58840304f55b81b61a775c0a5d4",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/313 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "from tqdm.auto import tqdm\n",
    "import warnings\n",
    "import pandas as pd\n",
    "import numpy as np\n",
    "\n",
    "warnings.filterwarnings(\"ignore\", category=UserWarning)\n",
    "\n",
    "# we will use batches of 64\n",
    "batch_size = 64\n",
    "data = []\n",
    "from collections import defaultdict\n",
    "\n",
    "# table_data = defaultdict(list)\n",
    "\n",
    "\n",
    "for i in tqdm(range(0, len(df), batch_size)):\n",
    "    # find end of batch\n",
    "    i_end = min(i + batch_size, len(df))\n",
    "    # extract batch\n",
    "    batch = df.iloc[i:i_end].copy()\n",
    "    # generate embeddings for batch\n",
    "    emb = retriever.encode(batch[\"title_text\"].tolist()).tolist()\n",
    "    # extract named entities from the batch\n",
    "    entities = extract_named_entities(batch[\"title_text\"].tolist())\n",
    "    # remove duplicate entities from each record\n",
    "    batch[\"named_entities\"] = [list(set(entity)) for entity in entities]\n",
    "    batch = batch.drop(\"title_text\", axis=1)\n",
    "    # get metadata\n",
    "    meta = batch.to_dict(orient=\"records\")\n",
    "    # create unique IDs\n",
    "    ids = [f\"{idx}\" for idx in range(i, i_end)]\n",
    "    # add all to upsert list\n",
    "    to_upsert = list(zip(ids, emb, meta, batch[\"named_entities\"]))\n",
    "    for id, emb, meta, entity in to_upsert:\n",
    "        temp = {}\n",
    "\n",
    "        temp[\"vector\"] = np.array(emb)\n",
    "        temp[\"metadata\"] = meta\n",
    "        temp[\"named_entities\"] = entity\n",
    "        data.append(temp)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "vRt4GuzYK8h2",
    "outputId": "9dae0ce7-9a18-49f3-a030-1fb1d5d5fedd"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "pyarrow.Table\n",
       "vector: fixed_size_list<item: float>[768]\n",
       "  child 0, item: float\n",
       "metadata: struct<authors: string, named_entities: list<item: string>, tags: string, text: string, timestamp: string, title: string, url: string>\n",
       "  child 0, authors: string\n",
       "  child 1, named_entities: list<item: string>\n",
       "      child 0, item: string\n",
       "  child 2, tags: string\n",
       "  child 3, text: string\n",
       "  child 4, timestamp: string\n",
       "  child 5, title: string\n",
       "  child 6, url: string\n",
       "named_entities: list<item: string>\n",
       "  child 0, item: string\n",
       "----\n",
       "vector: [[[-0.009049614,0.10612086,-0.027753588,0.07209486,-0.032509252,...,0.11016317,-0.013526588,-0.0046699173,0.035262693,-0.051537305],[-0.04104508,-0.049538508,-0.026324937,0.019106576,-0.017135208,...,-0.050371084,-0.058374014,0.014137886,-0.046907514,-0.012160475],[0.0068191127,0.05442695,0.0059523294,-0.0272331,0.05366467,...,0.05989369,-0.02457071,-0.01919812,0.059475537,-0.040533062],[-0.074406564,0.06398625,-0.0032167286,0.0006136286,-0.04038913,...,-0.0035826706,0.017684124,-0.039035726,-0.012032948,-0.0018114883],[0.0030152015,0.024747377,-0.009255136,0.00078397495,0.01981007,...,-0.03834965,-0.021975843,-0.029613132,-0.0023647717,-0.0017810423]]]\n",
       "metadata: [\n",
       "  -- is_valid: all not null\n",
       "  -- child 0 type: string\n",
       "[\"[]\",\"['Tomas Rasymas']\",\"['Yourself Quotes']\",\"['Matt Charnock']\",\"['Sip Wines']\"]\n",
       "  -- child 1 type: list<item: string>\n",
       "[[\"Data\",\"Christmas\",\"America\",\"Anonymous\"],[\"MicroPython\",\"Unsplash\",\"ESP32 Board\",\"Light Switch\",\"Jaye Haych\"],[\"Hope\",\"Keep Going\",\"Walker Morgan\"],[\"David King\",\"Martian\",\"AQI\",\"California\",\"SFGate\",\"Bay Area\",\"NWS\",\"Bay Area Skies\",\"National Weather Service\"],[\"ABC\",\"Julia DiPrete\",\"Jackson 5\",\"Patton Valley Vineyard\"]]\n",
       "  -- child 2 type: string\n",
       "[\"['Data Science', 'Big Data', 'Dataops', 'Analytics', 'DevOps']\",\"['Programming', 'Python', 'Software Development', 'Hardware', 'Automation']\",\"['Quotes']\",\"['Bay Area', 'San Francisco', 'California', 'Wildfires', 'Environment']\",\"['Wine Tasting', 'Sustainability', 'Wine']\"]\n",
       "  -- child 3 type: string\n",
       "[\"by Anonymous\n",
       "\n",
       "The door sprung open and our three little ones joyfully leaped onto the bed, waking my wife and I from peaceful dreams of slumbering sugarplum fairies. Oh Gosh, it was 6:03am, but who could blame them. The kids wait for Christmas day all year.\n",
       "\n",
       "I sat up as giggles thrashed with reckless excitement on the bed. My feet found slippers and I lumbered down the hall in the direction of the coffee maker. I could hear the tearing of wrapping paper and whoops of joy coming from the living-room tree.\n",
       "\n",
       "This scene was playing out in countless ways in households all across America. Christmas is a time for family, friends and a relaxing respite from the busy work calendar. Unfortunately, it was all about to come to an abrupt end as my cell phone interrupted our family time with an impatient buzz. “Who could be calling me today?”\n",
       "\n",
       "I should have powered it off and tossed the damn thing into a snowbank, but as the manager of a data team in a global e-commerce company, I was used to taking\",\"A story about how I escaped the boring task that I was doing twice a day by automating it!\n",
       "\n",
       "Photo by Jaye Haych on Unsplash\n",
       "\n",
       "A few weeks ago, I mounted outdoor lights on my house. It was a happy moment! 🥳 They added cosiness to my place and of course, made my yard lighter in the middle of the night.\n",
       "\n",
       "But at the time, I did not think about the switch too much. I installed a simple manual switch 🤦‍♂️ It seemed like no big deal, turning lights on for the night and turning them off in the morning. But after a few weeks, I had changed my mind!\n",
       "\n",
       "I started to think about how I could tackle this problem and I came up with a few ideas:\n",
       "\n",
       "Buy an automated light switch, that can be programmed to turn on and off the lights by the time of the day, or maybe over WiFi or something similar. More money is spent but the effect is guaranteed but I need to reinstall the switch.\n",
       "\n",
       "Another option is to do everything by myself. I am going to learn few new things, everything gonna cost less and I think I will n\",\"It’s a very thrilling thing to achieve a goal. If you’re battling with the idea of giving up, remember that if you give up, you could be giving away a wonderful thing — your best future. You may be fearful and uncertain of the future, but it’s crucial that you drive to keep going. After all, if you quit now, you’ll never know how close you were to succeeding. There are many things that can help you overcome your anxiety and fear.\n",
       "\n",
       "If at any point you think you are going to fall or lose hope, here are some Keep Going Quotes that will help empower you to never lose hope.\n",
       "\n",
       "If you have gone through something traumatic and you don’t feel like living anymore, don’t lose hope. There are many people who have faced major accident and burn injuries and still are doing great in life. For example, the national burn injury firm, Walker Morgan, has provided legal services for persons who have suffered severe burn injuries, helping them to get compensation for their medical expenses.\n",
       "\n",
       "There are many \",\"Bay Area cities are contending with some of the worst air quality in the world. We’ve been grappling with these dangerously toxic atmospheres for weeks now. Ironically enough: Last Wednesday’s Martian landscape was one of the few “moderate level” air quality days; Thursday through Sunday, however, saw record-breaking AQI recordings.\n",
       "\n",
       "Therefore, the question on everyone’s mind: When will Bay Area skies finally turn blue again and boast healthy air conditions? Well, the answer’s somewhat muddied, but clear, deep-azure overhead conditions are on the horizon — eventually.\n",
       "\n",
       "The National Weather Service (NWS) has tweeted a series of animated models over the past few days that suggest seaward gusts will inevitably sweep away our hazy skies by the middle of this week or possibly a little later than that.\n",
       "\n",
       "NWS meteorologist David King told SFGate on Sunday that the weather system is slated to blow its way to California’s coastline this week. The system is expected to carry with it increased win\",\"By Julia DiPrete\n",
       "\n",
       "(according to the Jackson 5, anyway)\n",
       "\n",
       "(Am I still allowed to like the Jackson 5? I’m so conflicted… Sigh.)\n",
       "\n",
       "Patton Valley Vineyard\n",
       "\n",
       "I’ll admit, sitting down to write a post about sustainability was daunting. For one thing, I knew very little about sustainability before writing this, so I had to learn. Then, once I started researching, I became slightly overwhelmed as to how I could distill all of it into a blog post of reasonable length. On top of all that, I had to figure out a way to present the distilled version of a serious and somewhat technical topic in a way that was attention-grabbing, entertaining, relatable — you know, all the things that make you actually want to read something. I hope you’ll read on, and as you do so, please keep in mind that I gave it my all.\n",
       "\n",
       "Here’s the thing: Sustainability is INCREDIBLY important… and it’s way more achievable than you might think.\n",
       "\n",
       "First things first: What is “sustainability” at its core?\n",
       "\n",
       "When I said I knew very littl\"]\n",
       "  -- child 4 type: string\n",
       "[\"2019-12-24 13:22:33.143000+00:00\",\"2021-09-14 07:20:52.342000+00:00\",\"2021-01-05 12:13:04.018000+00:00\",\"2020-09-15 22:38:33.924000+00:00\",\"2021-03-02 23:39:49.948000+00:00\"]\n",
       "  -- child 5 type: string\n",
       "[\"How the Data Stole Christmas\",\"Automating Light Switch using the ESP32 Board and MicroPython\",\"Keep Going Quotes Sayings for When Hope is Lost\",\"When Will the Smoke Clear From Bay Area Skies?\",\"The ABC’s of Sustainability… easy as 1, 2, 3\"]\n",
       "  -- child 6 type: string\n",
       "[\"https://medium.com/data-ops/how-the-data-stole-christmas-78454531d0a8\",\"https://python.plainenglish.io/automating-light-switch-using-esp32-board-and-micropython-502622f54242\",\"https://medium.com/@yourselfquotes/keep-going-quotes-sayings-for-when-hope-is-lost-92ad02ebc187\",\"https://thebolditalic.com/when-will-the-smoke-clear-from-bay-area-skies-f88d18ce94f5\",\"https://medium.com/sipwines/the-abcs-of-sustainability-easy-as-1-2-3-1a89e46c0f8a\"]]\n",
       "named_entities: [[[\"Data\",\"Christmas\",\"America\",\"Anonymous\"],[\"MicroPython\",\"Unsplash\",\"ESP32 Board\",\"Light Switch\",\"Jaye Haych\"],[\"Hope\",\"Keep Going\",\"Walker Morgan\"],[\"David King\",\"Martian\",\"AQI\",\"California\",\"SFGate\",\"Bay Area\",\"NWS\",\"Bay Area Skies\",\"National Weather Service\"],[\"ABC\",\"Julia DiPrete\",\"Jackson 5\",\"Patton Valley Vineyard\"]]]"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# create table using above data\n",
    "tbl = db.create_table(\"tw\", data)\n",
    "# check table\n",
    "tbl.head()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "uyZ5-INNxP19"
   },
   "source": [
    "# Quering"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {
    "id": "swjfItXsxRhv"
   },
   "outputs": [],
   "source": [
    "from pprint import pprint\n",
    "\n",
    "\n",
    "def search_lancedb(query):\n",
    "    # extract named entities from the query\n",
    "    ne = extract_named_entities([query])[0]\n",
    "    # create embeddings for the query\n",
    "    xq = retriever.encode(query).tolist()\n",
    "    # query the lancedb table while applying named entity filter\n",
    "    xc = tbl.search(xq).to_list()\n",
    "    # extract article titles from the search result\n",
    "    r = [\n",
    "        x[\"metadata\"][\"title\"]\n",
    "        for x in xc\n",
    "        for i in x[\"metadata\"][\"named_entities\"]\n",
    "        if i in ne\n",
    "    ]\n",
    "    return pprint({\"Extracted Named Entities\": ne, \"Result\": r})"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "X9XLV9hfxzkL"
   },
   "source": [
    "Now lets try quering"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "Q5gaOdIvxyo4",
    "outputId": "eb599773-5ff5-481d-a483-63d13ac888b6"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'Extracted Named Entities': ['Data'],\n",
      " 'Result': ['Data Science is all about making the right choices']}\n"
     ]
    }
   ],
   "source": [
    "query = \"How Data is changing world?\"\n",
    "search_lancedb(query)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "AI_-ALvjyDm2",
    "outputId": "079ac857-de73-4448-cc5e-97ff425c6767"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'Extracted Named Entities': ['SpaceX', 'Mars'],\n",
      " 'Result': ['Mars Habitat: NASA 3D Printed Habitat Challenge',\n",
      "            'Reusable rockets and the robots at sea: The SpaceX story',\n",
      "            'Reusable rockets and the robots at sea: The SpaceX story',\n",
      "            'Colonising Planets Beyond Mars',\n",
      "            'Colonising Planets Beyond Mars',\n",
      "            'Musk Explained: The Musk approach to marketing',\n",
      "            'How We’ll Access the Water on Mars',\n",
      "            'Chasing Immortality',\n",
      "            'Mission Possible: How Space Exploration Can Deliver Sustainable '\n",
      "            'Development']}\n"
     ]
    }
   ],
   "source": [
    "query = \"Why does SpaceX want to build a city on Mars?\"\n",
    "search_lancedb(query)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "aC0D5gKroiJt"
   },
   "source": [
    "These all look like great results, making the most of LanceDB advanced vector search capabilities"
   ]
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "gpuType": "T4",
   "provenance": []
  },
  "kernelspec": {
   "display_name": "Python 3",
   "name": "python3"
  },
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}
